#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
生成专业学术风格的组件频率热力图

功能：
1. 读取parquet文件中的验证组件信息
2. 统计每个组件在verification_components中的出现次数
3. 生成高质量的学术风格热力图
4. 清晰区分服务、容器、节点三种类型
5. 分别生成Top-1和非Top-1样本的热力图

作者：AI Assistant
日期：2025-11-03
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from collections import Counter, defaultdict
from typing import Dict, List, Tuple, Any
import argparse
from pathlib import Path

# 设置matplotlib支持中文
plt.rcParams['font.sans-serif'] = ['Arial Unicode MS', 'DejaVu Sans', 'SimHei']
plt.rcParams['axes.unicode_minus'] = False

# 学术出版质量设置
plt.rcParams['figure.dpi'] = 300
plt.rcParams['savefig.dpi'] = 300
plt.rcParams['figure.figsize'] = (16, 10)


def get_component_type(component: str) -> str:
    """
    判断组件类型
    
    参数:
        component: 组件名称
        
    返回:
        'service', 'container', 或 'node'
    """
    if component.startswith('node-'):
        return 'node'
    elif '-' in component:
        return 'container'
    else:
        return 'service'


def get_service_from_container(container: str) -> str:
    """
    从容器名称提取服务名
    
    参数:
        container: 容器名称，如 "paymentservice-0"
        
    返回:
        服务名称，如 "paymentservice"
    """
    if '-' in container and not container.startswith('node-'):
        # 处理类似 "paymentservice-0" 或 "paymentservice2-0" 的情况
        parts = container.split('-')
        if len(parts) >= 2:
            # 去掉最后的数字部分
            return '-'.join(parts[:-1]) if parts[-1].isdigit() else container
    return container


def analyze_components_frequency(df: pd.DataFrame) -> Tuple[Dict, Dict]:
    """
    分析verification_components的频率分布
    
    参数:
        df: parquet数据的DataFrame
        
    返回:
        (top1_freq, non_top1_freq) 两个频率计数器
    """
    top1_components = Counter()
    non_top1_components = Counter()
    
    top1_count = 0
    non_top1_count = 0
    
    for idx, row in df.iterrows():
        record = row.to_dict()
        
        try:
            info = record.get('info')
            if not info or not isinstance(info, dict):
                continue
            
            # 获取排名
            metrics = info.get('metrics', {})
            rank = metrics.get('AIOps-22/root_cause_rank', 999)
            
            # 获取verification_components
            env_info = info.get('env_info', {})
            print_result_data = env_info.get('print_result_data', {})
            verification_components = print_result_data.get('verification_components', [])
            
            # 处理numpy数组
            if isinstance(verification_components, np.ndarray):
                verification_components = verification_components.tolist()
            
            if not isinstance(verification_components, list):
                continue
            
            # 统计组件出现次数
            if rank == 1:
                top1_count += 1
                for comp in verification_components:
                    top1_components[str(comp)] += 1
            elif rank < 999:
                non_top1_count += 1
                for comp in verification_components:
                    non_top1_components[str(comp)] += 1
        
        except Exception as e:
            print(f"处理记录 {idx} 时出错: {e}")
            continue
    
    print(f"\n样本统计:")
    print(f"  Top-1样本数: {top1_count}")
    print(f"  非Top-1样本数: {non_top1_count}")
    
    return dict(top1_components), dict(non_top1_components)


def prepare_heatmap_data(component_freq: Dict[str, int]) -> Tuple[pd.DataFrame, List, List, List]:
    """
    准备热力图数据
    
    参数:
        component_freq: 组件频率字典
        
    返回:
        (data_df, services, containers, nodes)
    """
    # 按类型分组
    services = []
    containers = []
    nodes = []
    
    for comp in sorted(component_freq.keys()):
        comp_type = get_component_type(comp)
        if comp_type == 'service':
            services.append(comp)
        elif comp_type == 'container':
            containers.append(comp)
        else:
            nodes.append(comp)
    
    # 按服务分组容器
    service_containers = defaultdict(list)
    for container in containers:
        service = get_service_from_container(container)
        service_containers[service].append(container)
    
    # 构建数据矩阵：行为组件，列为频率
    all_components = []
    frequencies = []
    types = []
    groups = []
    
    # 添加服务
    for service in sorted(services):
        all_components.append(service)
        frequencies.append(component_freq.get(service, 0))
        types.append('Service')
        groups.append(service)
    
    # 按服务分组添加容器
    for service in sorted(set(get_service_from_container(c) for c in containers)):
        service_conts = sorted([c for c in containers if get_service_from_container(c) == service])
        for container in service_conts:
            all_components.append(container)
            frequencies.append(component_freq.get(container, 0))
            types.append('Container')
            groups.append(service)
    
    # 添加节点
    for node in sorted(nodes):
        all_components.append(node)
        frequencies.append(component_freq.get(node, 0))
        types.append('Node')
        groups.append('Infrastructure')
    
    # 创建DataFrame
    data_df = pd.DataFrame({
        'Component': all_components,
        'Frequency': frequencies,
        'Type': types,
        'Group': groups
    })
    
    return data_df, services, containers, nodes


def plot_professional_heatmap(
    data_df: pd.DataFrame,
    title: str,
    output_file: str,
    total_samples: int
):
    """
    绘制专业学术风格的热力图（按百分比着色）
    
    参数:
        data_df: 包含组件和频率的DataFrame
        title: 图表标题
        output_file: 输出文件路径
        total_samples: 总样本数（用于显示，不用于着色）
    """
    # 创建图形
    fig, ax = plt.subplots(figsize=(18, max(12, len(data_df) * 0.4)))
    
    # 准备数据矩阵（单列）
    matrix = data_df['Frequency'].values.reshape(-1, 1)
    
    # 计算百分比（用于着色）- 关键：占所有组件总出现次数的比例
    total_occurrences = matrix.sum()  # 所有组件的总出现次数
    percentages = (matrix / total_occurrences * 100) if total_occurrences > 0 else matrix
    
    # 创建热力图
    # 使用YlOrRd配色方案（黄-橙-红）更专业
    cmap = sns.color_palette("YlOrRd", as_cmap=True)
    
    # 绘制热力图 - 关键：使用百分比矩阵来着色，但显示绝对频率
    sns.heatmap(
        percentages,  # 使用百分比来决定颜色
        annot=matrix.astype(int),  # 但显示绝对频率数值
        fmt='d',
        cmap=cmap,
        cbar_kws={'label': '% of Total Occurrences'},  # 颜色条显示占总出现次数的百分比
        linewidths=0.5,
        linecolor='gray',
        ax=ax,
        vmin=0,
        vmax=percentages.max() if percentages.max() > 0 else 100,  # 动态最大值
        square=False,
        xticklabels=['Verification Frequency'],
        yticklabels=data_df['Component'].values
    )
    
    # 添加百分比标注
    for i, (freq, pct) in enumerate(zip(matrix.flatten(), percentages.flatten())):
        if freq > 0:
            ax.text(0.5, i + 0.7, f'({pct:.1f}%)', 
                   ha='center', va='center', fontsize=8, color='darkblue')
    
    # 根据组件类型设置背景颜色
    for i, row in data_df.iterrows():
        if row['Type'] == 'Service':
            ax.add_patch(plt.Rectangle((-0.5, i), 0.05, 1, 
                                       fill=True, color='#2196F3', alpha=0.3, 
                                       transform=ax.transData, clip_on=False))
        elif row['Type'] == 'Container':
            ax.add_patch(plt.Rectangle((-0.5, i), 0.05, 1, 
                                       fill=True, color='#FF9800', alpha=0.3, 
                                       transform=ax.transData, clip_on=False))
        else:  # Node
            ax.add_patch(plt.Rectangle((-0.5, i), 0.05, 1, 
                                       fill=True, color='#9C27B0', alpha=0.3, 
                                       transform=ax.transData, clip_on=False))
    
    # 设置标题和标签
    ax.set_title(title, fontsize=16, fontweight='bold', pad=20)
    ax.set_xlabel('Metric', fontsize=12, fontweight='bold')
    ax.set_ylabel('Component', fontsize=12, fontweight='bold')
    
    # 调整y轴标签
    ax.set_yticklabels(ax.get_yticklabels(), fontsize=9, rotation=0)
    
    # 添加图例
    from matplotlib.patches import Patch
    legend_elements = [
        Patch(facecolor='#2196F3', alpha=0.3, label='Service'),
        Patch(facecolor='#FF9800', alpha=0.3, label='Container'),
        Patch(facecolor='#9C27B0', alpha=0.3, label='Node')
    ]
    ax.legend(handles=legend_elements, loc='upper right', 
             bbox_to_anchor=(1.15, 1), fontsize=10, title='Component Type')
    
    # 添加统计信息
    max_freq = data_df['Frequency'].max()
    total_freq = data_df['Frequency'].sum()
    non_zero = (data_df['Frequency'] > 0).sum()
    max_pct = percentages.max()
    
    stats_text = f"Total Samples: {total_samples}\n"
    stats_text += f"Total Components: {len(data_df)}\n"
    stats_text += f"Active Components: {non_zero}\n"
    stats_text += f"Max Frequency: {int(max_freq)}\n"
    stats_text += f"Total Occurrences: {int(total_freq)}\n"
    stats_text += f"Max Percentage: {max_pct:.1f}%"
    
    ax.text(1.18, 0.5, stats_text, transform=ax.transAxes,
           fontsize=10, verticalalignment='center',
           bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.3))
    
    # 调整布局
    plt.tight_layout()
    
    # 保存图形
    plt.savefig(output_file, dpi=300, bbox_inches='tight')
    print(f"✅ 已保存到: {output_file}")
    
    plt.close()


def plot_grouped_bar_chart(
    data_df: pd.DataFrame,
    title: str,
    output_file: str,
    total_samples: int
):
    """
    绘制分组柱状图（按服务分组）
    
    参数:
        data_df: 包含组件和频率的DataFrame
        title: 图表标题
        output_file: 输出文件路径
        total_samples: 总样本数
    """
    # 创建图形
    fig, ax = plt.subplots(figsize=(16, 8))
    
    # 准备数据
    df_sorted = data_df.sort_values('Frequency', ascending=False)
    
    # 根据类型设置颜色
    colors = []
    for comp_type in df_sorted['Type']:
        if comp_type == 'Service':
            colors.append('#2196F3')
        elif comp_type == 'Container':
            colors.append('#FF9800')
        else:
            colors.append('#9C27B0')
    
    # 绘制柱状图
    bars = ax.bar(range(len(df_sorted)), df_sorted['Frequency'], color=colors, alpha=0.7, edgecolor='black')
    
    # 添加数值标签
    for i, (bar, freq) in enumerate(zip(bars, df_sorted['Frequency'])):
        if freq > 0:
            height = bar.get_height()
            pct = (freq / total_samples * 100) if total_samples > 0 else 0
            ax.text(bar.get_x() + bar.get_width()/2., height,
                   f'{int(freq)}\n({pct:.1f}%)',
                   ha='center', va='bottom', fontsize=8)
    
    # 设置x轴标签
    ax.set_xticks(range(len(df_sorted)))
    ax.set_xticklabels(df_sorted['Component'], rotation=45, ha='right', fontsize=9)
    
    # 设置标题和标签
    ax.set_title(title, fontsize=16, fontweight='bold', pad=20)
    ax.set_xlabel('Component', fontsize=12, fontweight='bold')
    ax.set_ylabel('Frequency', fontsize=12, fontweight='bold')
    
    # 添加网格
    ax.grid(axis='y', alpha=0.3, linestyle='--')
    ax.set_axisbelow(True)
    
    # 添加图例
    from matplotlib.patches import Patch
    legend_elements = [
        Patch(facecolor='#2196F3', alpha=0.7, label='Service'),
        Patch(facecolor='#FF9800', alpha=0.7, label='Container'),
        Patch(facecolor='#9C27B0', alpha=0.7, label='Node')
    ]
    ax.legend(handles=legend_elements, loc='upper right', fontsize=10)
    
    # 调整布局
    plt.tight_layout()
    
    # 保存图形
    plt.savefig(output_file, dpi=300, bbox_inches='tight')
    print(f"✅ 已保存到: {output_file}")
    
    plt.close()


def plot_comparison_heatmap(
    top1_df: pd.DataFrame,
    non_top1_df: pd.DataFrame,
    output_file: str,
    top1_samples: int,
    non_top1_samples: int
):
    """
    绘制Top-1和非Top-1的对比热力图（按百分比着色）
    
    参数:
        top1_df: Top-1的DataFrame
        non_top1_df: 非Top-1的DataFrame
        output_file: 输出文件路径
        top1_samples: Top-1样本数
        non_top1_samples: 非Top-1样本数
    """
    # 合并所有组件
    all_components = sorted(set(top1_df['Component'].tolist() + non_top1_df['Component'].tolist()))
    
    # 创建对比矩阵（频率和百分比）
    freq_matrix = []
    pct_matrix = []
    
    # 计算总出现次数（每列独立计算）
    top1_total_occurrences = top1_df['Frequency'].sum()
    non_top1_total_occurrences = non_top1_df['Frequency'].sum()
    
    for comp in all_components:
        top1_freq = top1_df[top1_df['Component'] == comp]['Frequency'].values
        top1_freq = top1_freq[0] if len(top1_freq) > 0 else 0
        # 占Top-1所有组件总出现次数的比例
        top1_pct = (top1_freq / top1_total_occurrences * 100) if top1_total_occurrences > 0 else 0
        
        non_top1_freq = non_top1_df[non_top1_df['Component'] == comp]['Frequency'].values
        non_top1_freq = non_top1_freq[0] if len(non_top1_freq) > 0 else 0
        # 占Non-Top-1所有组件总出现次数的比例
        non_top1_pct = (non_top1_freq / non_top1_total_occurrences * 100) if non_top1_total_occurrences > 0 else 0
        
        freq_matrix.append([top1_freq, non_top1_freq])
        pct_matrix.append([top1_pct, non_top1_pct])
    
    freq_matrix = np.array(freq_matrix)
    pct_matrix = np.array(pct_matrix)
    
    # 创建图形
    fig, ax = plt.subplots(figsize=(12, max(12, len(all_components) * 0.4)))
    
    # 绘制热力图 - 使用百分比着色，显示频率
    sns.heatmap(
        pct_matrix,  # 使用百分比来决定颜色
        annot=freq_matrix.astype(int),  # 显示绝对频率
        fmt='d',
        cmap='YlOrRd',
        cbar_kws={'label': '% of Total Occurrences'},  # 占总出现次数的百分比
        linewidths=0.5,
        linecolor='gray',
        ax=ax,
        vmin=0,
        vmax=pct_matrix.max() if pct_matrix.max() > 0 else 100,  # 动态最大值
        xticklabels=[f'Top-1 (n={top1_samples})', f'Non-Top-1 (n={non_top1_samples})'],
        yticklabels=all_components
    )
    
    # 添加百分比标注
    for i in range(len(all_components)):
        for j in range(2):
            freq = freq_matrix[i, j]
            pct = pct_matrix[i, j]
            if freq > 0:
                ax.text(j + 0.5, i + 0.7, f'({pct:.1f}%)',
                       ha='center', va='center', fontsize=8, color='darkblue')
    
    # 设置标题
    ax.set_title('Comparison: Top-1 vs Non-Top-1 Component Frequency', 
                fontsize=16, fontweight='bold', pad=20)
    ax.set_ylabel('Component', fontsize=12, fontweight='bold')
    
    # 调整y轴标签
    ax.set_yticklabels(ax.get_yticklabels(), fontsize=9, rotation=0)
    
    # 添加统计信息
    stats_text = f"Top-1 Samples: {top1_samples}\n"
    stats_text += f"Non-Top-1 Samples: {non_top1_samples}\n"
    stats_text += f"Total Components: {len(all_components)}\n"
    stats_text += f"\n"
    stats_text += f"Top-1 Total Occur.: {int(top1_total_occurrences)}\n"
    stats_text += f"Non-Top-1 Total Occur.: {int(non_top1_total_occurrences)}"
    
    ax.text(1.15, 0.5, stats_text, transform=ax.transAxes,
           fontsize=10, verticalalignment='center',
           bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.3))
    
    # 调整布局
    plt.tight_layout()
    
    # 保存图形
    plt.savefig(output_file, dpi=300, bbox_inches='tight')
    print(f"✅ 已保存到: {output_file}")
    
    plt.close()


def main():
    """主函数"""
    parser = argparse.ArgumentParser(
        description='生成专业学术风格的组件频率热力图'
    )
    parser.add_argument(
        'parquet_file',
        type=str,
        help='输入的parquet文件路径'
    )
    parser.add_argument(
        '--output-dir',
        type=str,
        default='.',
        help='输出目录路径'
    )
    
    args = parser.parse_args()
    
    print(f"正在读取文件: {args.parquet_file}")
    df = pd.read_parquet(args.parquet_file)
    total_records = len(df)
    print(f"总记录数: {total_records}")
    
    # 分析验证组件频率
    print(f"\n正在分析verification_components频率...")
    top1_freq, non_top1_freq = analyze_components_frequency(df)
    
    top1_samples = sum(1 for row in df.iterrows() 
                      if row[1].to_dict().get('info', {}).get('metrics', {}).get('AIOps-22/root_cause_rank', 999) == 1)
    non_top1_samples = total_records - top1_samples
    
    print(f"✅ 分析完成")
    
    # 创建输出目录
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    
    # 准备Top-1数据
    print(f"\n正在生成Top-1样本热力图...")
    top1_df, _, _, _ = prepare_heatmap_data(top1_freq)
    
    # 绘制Top-1热力图
    plot_professional_heatmap(
        top1_df,
        'Top-1 Samples: Component Verification Frequency Heatmap',
        str(output_dir / 'heatmap_top1_professional.png'),
        top1_samples
    )
    
    # 绘制Top-1柱状图
    plot_grouped_bar_chart(
        top1_df,
        'Top-1 Samples: Component Verification Frequency',
        str(output_dir / 'barchart_top1.png'),
        top1_samples
    )
    
    # 准备非Top-1数据
    if non_top1_samples > 0:
        print(f"\n正在生成非Top-1样本热力图...")
        non_top1_df, _, _, _ = prepare_heatmap_data(non_top1_freq)
        
        # 绘制非Top-1热力图
        plot_professional_heatmap(
            non_top1_df,
            'Non-Top-1 Samples: Component Verification Frequency Heatmap',
            str(output_dir / 'heatmap_non_top1_professional.png'),
            non_top1_samples
        )
        
        # 绘制非Top-1柱状图
        plot_grouped_bar_chart(
            non_top1_df,
            'Non-Top-1 Samples: Component Verification Frequency',
            str(output_dir / 'barchart_non_top1.png'),
            non_top1_samples
        )
        
        # 绘制对比图
        print(f"\n正在生成对比热力图...")
        plot_comparison_heatmap(
            top1_df,
            non_top1_df,
            str(output_dir / 'heatmap_comparison.png'),
            top1_samples,
            non_top1_samples
        )
    
    # 生成统计摘要
    print(f"\n" + "="*80)
    print("统计摘要")
    print("="*80)
    
    print(f"\n【Top-1样本统计】")
    print(f"  样本数: {top1_samples}")
    print(f"  不同组件数: {len(top1_df)}")
    print(f"  活跃组件数: {(top1_df['Frequency'] > 0).sum()}")
    print(f"\n  Top-10 最常出现的组件:")
    for idx, row in top1_df.nlargest(10, 'Frequency').iterrows():
        pct = (row['Frequency'] / top1_samples * 100) if top1_samples > 0 else 0
        print(f"    {row['Component']:30s} ({row['Type']:10s}): {row['Frequency']:3d}次 ({pct:5.1f}%)")
    
    if non_top1_samples > 0:
        print(f"\n【非Top-1样本统计】")
        print(f"  样本数: {non_top1_samples}")
        print(f"  不同组件数: {len(non_top1_df)}")
        print(f"  活跃组件数: {(non_top1_df['Frequency'] > 0).sum()}")
        print(f"\n  Top-10 最常出现的组件:")
        for idx, row in non_top1_df.nlargest(10, 'Frequency').iterrows():
            pct = (row['Frequency'] / non_top1_samples * 100) if non_top1_samples > 0 else 0
            print(f"    {row['Component']:30s} ({row['Type']:10s}): {row['Frequency']:3d}次 ({pct:5.1f}%)")
    
    print(f"\n✅ 所有图表已生成")
    print(f"输出目录: {output_dir}")
    print(f"\n生成的文件:")
    print(f"  - heatmap_top1_professional.png     (Top-1热力图)")
    print(f"  - barchart_top1.png                 (Top-1柱状图)")
    if non_top1_samples > 0:
        print(f"  - heatmap_non_top1_professional.png (非Top-1热力图)")
        print(f"  - barchart_non_top1.png             (非Top-1柱状图)")
        print(f"  - heatmap_comparison.png            (对比热力图)")
    
    return 0


if __name__ == "__main__":
    exit_code = main()
    exit(exit_code)

